﻿using Fst;
using Fst.Operations;
using Fst.SemiRing;
using Fst.Utils;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;

namespace Linguist.G2p
{
    /// <summary>
    /// The grapheme-to-phoneme (g2p) decoder
    /// 
    /// @author John Salatas <jsalatas@users.sourceforge.net>
    /// </summary>
    public class G2PConverter
    {
        // epsilon symbol
        String eps = "<eps>";

        // end sequence symbol
        String se = "</s>";

        // begin sequence symbol
        String sb = "<s>";

        // skip symbol
        String skip = "_";

        // separator symbol
        String tie = "|";

        // set containing sequences to ignore
        HashSet<String> skipSeqs = new HashSet<String>();

        // clusters
        List<List<String>> clusters = null;

        // the g2p model
        ImmutableFst g2pmodel;

        // fst containing the epsilon filter for the compose operation
        Fst.Fst epsilonFilter;

        /**
        /// Create a decoder by loading the serialized model from a specified URL
        /// 
        /// @param g2pModelUrl
        ///            the URL of the serialized model
        /// @throws IOException
        /// @throws ClassNotFoundException 
         */
        public G2PConverter(String g2pModelUrl)  
        {
            try {
                g2pmodel = ImmutableFst.loadModel(g2pModelUrl);
            } 
            catch (Exception e) 
            {
                throw new IOException("Failed to load the model from " + g2pModelUrl, e);
            }
            init();
        }

        ///**
        ///// Create a decoder by loading the serialized model from a specified
        ///// filename
        ///// 
        ///// @param g2pmodel_file
        /////            the filename of the serialized model
        // */
        //public G2PConverter(String g2pmodel_file) 
        //{
        //    g2pmodel = ImmutableFst.loadModel(g2pmodel_file);
        //    init();
        //}

        /**
        /// Initialize the decoder
         */
        private void init() {
            skipSeqs.Add(eps);
            skipSeqs.Add(sb);
            skipSeqs.Add(se);
            skipSeqs.Add(skip);
            skipSeqs.Add("-");
            // keep an augmented copy (for compose)
            Compose.augment(0, g2pmodel, g2pmodel.getSemiring());
            ArcSort.apply(g2pmodel, new ILabelCompare());

            String[] isyms = g2pmodel.getIsyms();

            loadClusters(isyms);

            // get epsilon filter for composition
            epsilonFilter = Compose.getFilter(g2pmodel.getIsyms(),g2pmodel.getSemiring());
            ArcSort.apply(epsilonFilter, new ILabelCompare());
        }

        /**
        /// Phoneticize a word
        /// 
        /// @param entry
        ///            the word to phoneticize transformed to an ArrayList of Strings
        ///            (each element hold a single character)
        /// @param nbest
        ///            the number of distinct pronunciations to return
        /// @return the pronunciation(s) of the input word
         */
        public List<Path> phoneticize(List<String> entry, int nbest) 
        {
            Fst.Fst efst = entryToFSA(entry);
            Semiring s = efst.getSemiring();
            Compose.augment(1, efst, s);
            ArcSort.apply(efst, new OLabelCompare());
            Fst.Fst result = Compose.compose(efst, epsilonFilter, s, true);
            ArcSort.apply(result, new OLabelCompare());
            result = Compose.compose(result, g2pmodel, s, true);
            Project.apply(result, ProjectType.OUTPUT);
            if (nbest == 1) {
                result = NShortestPaths.get(result, 1, false);
            } else {
                // Requesting 10 times more best paths than what was asking
                // as there might be several paths resolving to same pronunciation
                // due to epsilon transitions.
                // I really hate cosmological constants :)
                result = NShortestPaths.get(result, nbest* 10, false);
            }
            // result = NShortestPaths.get(result, nbest, false);
            result = RmEpsilon.get(result);
            List<Path> paths = findAllPaths(result, nbest, skipSeqs,tie);

            return paths;
        }

        /**
        /// Phoneticize a word
        /// 
        /// @param entry
        ///            the word to phoneticize
        /// @param nbest
        ///            the number of distinct pronunciations to return
        /// @return the pronunciation(s) of the input word
         */
        public List<Path> phoneticize(String word, int nbest) 
        {
            List<String> entry = new List<String>(word.Length);
            for (int i = 0; i < word.Length; i++) {
                String ch = word.Substring(i, i + 1);
                if (Utils.getIndex(g2pmodel.getIsyms(), ch) >= 0) 
                {
                    entry.Add(ch);
                }
            }
            return phoneticize(entry, nbest);
        }

        /**
        /// Transforms an input spelling/pronunciation into an equivalent FSA, adding
        /// extra arcs as needed to accommodate clusters.
        /// 
        /// @param entry
        ///            the input vector
        /// @return the created fst
         */
        private Fst.Fst entryToFSA(List<String> entry) 
        {
            TropicalSemiring ts = new TropicalSemiring();
            Fst.Fst efst = new Fst.Fst(ts);

            State s = new State(ts.zero());
            efst.addState(s);
            efst.setStart(s);

            // Build the basic FSA
            for (int i = 0; i < entry.Count + 1; i++) {
                s = new State(ts.zero());
                efst.addState(s);
                if (i >= 1) {
                    int symIndex = Utils.getIndex(g2pmodel.getIsyms(),
                            entry[i - 1]);
                    efst.getState(i).addArc(new Arc(symIndex, symIndex, 0.0f, s));
                } else if (i == 0) {
                    int symIndex = Utils.getIndex(g2pmodel.getIsyms(), sb);
                    efst.getStart().addArc(new Arc(symIndex, symIndex, 0.0f, s));
                }

                if (i == entry.Count) {
                    State s1 = new State(ts.zero());
                    efst.addState(s1);
                    int symIndex = Utils.getIndex(g2pmodel.getIsyms(), se);
                    s.addArc(new Arc(symIndex, symIndex, 0.0f, s1));
                    s1.setFinalWeight(0.0f);
                }
            }

            // Add any cluster arcs
            for (int value = 0; value < clusters.Count; value++) {
                List<String> cluster = clusters[value];
                if (cluster != null) {
                    int start = 0;
                    int k = 0;
                    while (k != -1) {
                        k = Utils.search(entry, cluster, start);
                        if (k != -1) {
                            State from = efst.getState(start + k + 1);
                            from.addArc(new Arc(value, value, 0.0f, efst
                                    .getState(start + k + cluster.Count + 1)));
                            start = start + k + cluster.Count;
                        }
                    }
                }
            }

            efst.setIsyms(g2pmodel.getIsyms());
            efst.setOsyms(g2pmodel.getIsyms());

            return efst;
        }

        /**
        /// Finds nbest paths in an Fst returned by NShortestPaths operation
        /// 
        /// @param fst
        ///            the input fst
        /// @param nbest
        ///            the number of paths to return
        /// @param skipSeqs
        ///            the sequences to ignore
        /// @param tie
        ///            the separator symbol
        /// @return the paths
         */
        private List<Path> findAllPaths(Fst.Fst fst, int nbest, HashSet<String> skipSeqs, String tie) 
        {
            Semiring semiring = fst.getSemiring();

            // ArrayList<Path> finalPaths = new ArrayList<Path>();
            Dictionary<String, Path> finalPaths = new Dictionary<String, Path>();
            Dictionary<State, Path> paths = new Dictionary<State, Path>();
            Queue<State> queue = new Queue<State>();
            Path p = new Path(fst.getSemiring());
            p.setCost(semiring.one());
            paths.Add(fst.getStart(), p);

            queue.Enqueue(fst.getStart());

            String[] osyms = fst.getOsyms();
            while (queue.Count!=0) 
            {
                State s = queue.Dequeue();
                Path currentPath = paths[s];

                if (s.getFinalWeight() != semiring.zero()) {
                    String pathString = currentPath.getPath().ToString();
                    if (finalPaths.ContainsKey(pathString)) 
                    {
                        // path already exist. update its cost
                        Path old = finalPaths[pathString];
                        if (old.getCost() > currentPath.getCost()) 
                        {
                            finalPaths.Add(pathString, currentPath);
                        }
                    } 
                    else 
                    {
                        finalPaths.Add(pathString, currentPath);
                    }
                }

                int numArcs = s.getNumArcs();
                for (int j = 0; j < numArcs; j++) 
                {
                    Arc a = s.getArc(j);
                    p = new Path(fst.getSemiring());
                    Path cur = paths[s];
                    p.setCost(cur.getCost());
                    p.setPath(cur.getPath().ToList());

                    String sym = osyms[a.getOlabel()];

                    String[] symsArray = sym.Split(new String[]{"\\" + tie}, StringSplitOptions.None);

                    for (int i = 0; i < symsArray.Length; i++) 
                    {
                        String phone = symsArray[i];
                        if (!skipSeqs.Contains(phone)) 
                        {
                            p.getPath().Add(phone);
                        }
                    }
                    p.setCost(semiring.times(p.getCost(), a.getWeight()));
                    State nextState = a.getNextState();
                    paths.Add(nextState, p);
                    if (!queue.Contains(nextState)) 
                    {
                        queue.Enqueue(nextState);
                    }
                }
            }

            List<Path> res = new List<Path>();
            foreach (Path path in finalPaths.Values) 
            {
                res.Add(path);
            }

            res.Sort(new PathComparator());
            int numPaths = res.Count;
            for (int i = nbest; i < numPaths; i++) 
            {
                res.RemoveAt(res.Count - 1);
            }

            return res;
        }

        /**
        /// Initialize clusters
         */
        private void loadClusters(String[] syms) 
        {
            clusters = new List<List<String>>();
            if (syms == null)
                return;
            for (int i = 0; i < syms.Length; i++) 
            {
                if (i < 2)
                {
                    clusters.Add(null);
                    continue;
                }
                String sym = syms[i];
                if (sym.Contains(tie)) 
                {
                    String[] split = sym.Split(new String[]{tie},StringSplitOptions.None);
                    List<String> clusterString = split.ToList();
                    clusters.Add(clusterString);
                }
            }
        }
    }
}
